import React, { useRef, useImperativeHandle, forwardRef } from 'react';
import DOMPurify from 'dompurify';

interface HighlightProps {
    children: React.ReactNode;
    highlightClassName?: string;
}

export interface HighlightHandle {
    highlight: (text: string) => void;
    removeHighlight: () => void;
}

const Highlight = forwardRef<HighlightHandle, HighlightProps>(({ children, highlightClassName }, ref) => {
    const textRef = useRef<HTMLDivElement>(null);

    useImperativeHandle(ref, () => ({
        highlight: (text: string) => {
            if (textRef.current) {
                highlightText(textRef.current, DOMPurify.sanitize(text));
            }
        },
        removeHighlight: () => {
            if (textRef.current) {
                removeHighlight(textRef.current);
            }
        }
    }));

    const highlightText = (node: Node, text: string) => {
        if (node.nodeType === Node.TEXT_NODE) {
            const index = (node.nodeValue as string).indexOf(text);
            if (index !== -1) {
                const span = document.createElement('span');
                if (highlightClassName) {
                    span.className = highlightClassName;
                } else {
                    span.style.color = 'red';
                }
                span.textContent = text;

                const before = (node as Text).splitText(index);
                before.parentNode?.insertBefore(span, before);
                (before.nodeValue as string) = (before.nodeValue as string)?.slice(text.length);
            }
        } else if (node.nodeType === Node.ELEMENT_NODE) {
            node.childNodes.forEach(child => highlightText(child, text));
        }
    };


    const removeHighlight = (node: Node) => {
        if (node.nodeType === Node.ELEMENT_NODE) {
            const spans = (node as Element).querySelectorAll(highlightClassName ? `.${highlightClassName}` : 'span[style="color: red;"]');
            spans.forEach(span => {
                const parent = span.parentNode;
                parent?.replaceChild(document.createTextNode(span.textContent as string), span);
                parent?.normalize();
            });
        }
    };

    return <div ref={textRef}>{children}</div>;
})

export default Highlight;
